import copy
from typing import Dict, Tuple, List

import iminuit
import pandas


class _ParamFitted:
    """
    参数拟合结果
    """

    def __init__(self, name: str = "", lower: float = None, upper: float = None, init: float = 0,
                 value: float = 0., error: float = 1.):
        self.name: str = name  # 参数名称
        self.lower: float | None = lower  # 下界
        self.upper: float | None = upper  # 上界
        self.init: float = init  # 初始值
        self.value: float = value  # 拟合值
        self.error: float = error  # 误差


class MultiParamProfiler:
    def __init__(self):
        self._param_fitted_dict: Dict[str, _ParamFitted] = dict()  # 参数拟合结果
        self._func = None  # 拟合函数

    def fit(
            self,
            df_x: pandas.DataFrame,
            df_y: pandas.DataFrame,
            df_param: pandas.DataFrame,
            func,
            y_col: str = "y",
            error_y_col: str = "error_y",
    ) -> pandas.DataFrame:
        """
        拟合
        :param df_x: 拟合的y=f(x)函数中的x值，其中列名和拟合函数中的 x_dict 对应的key值对应
        :param df_y: 拟合的y=f(x)函数中的y值和y的误差，其中列名和参数 y_col、error_y_col保持一致，默认为“y”、“error_y”。error_y非0
        :param df_param: 拟合函数中的参数信息，列名为name、init、lower、upper，分别为参数名、初始值、下界、上界
        :param func: 拟合函数（也被称机理函数、参数化公式等），输入参数x_dict为x值字典、param_dict为参数值字典，输出为float对应y值。
        :param y_col: df_y中y值的列名，默认为 y
        :param error_y_col: df_y中y值的误差的列名，默认为 error_y
        :return: 参数拟合结果，其列名 name、init、lower、upper、value、error，分别对应 参数名、初始值、下界、上界、拟合结果、拟合误差
        """
        return self._fit_inner(df_x=df_x, df_y=df_y, df_param=df_param, func=func, y_col=y_col, error_y_col=error_y_col)

    def predict(
            self,
            x_dict: Dict[str, float],
            func_error=None,
    ) -> Tuple[float, float | None]:
        """
        推理函数
        :param x_dict: 拟合的y=f(x)函数中的x值，其中的key值和和拟合函数中df_x的列名对应
        :param func_error: 误差计算函数，输入参数 x_dict 为x值的字典、param_dict为参数值字典，输出 float 为y值误差。
        :return: y值和其误差。当误差计算函数为None时，结果的误差值为None。
        """
        return self._predict_inner(x_dict=x_dict, func_error=func_error)

    def _fit_inner(
            self,
            df_x: pandas.DataFrame,
            df_y: pandas.DataFrame,
            df_param: pandas.DataFrame,
            func,
            y_col: str = "y",
            error_y_col: str = "error_y",
    ) -> pandas.DataFrame:
        """
        拟合
        :param df_x: 各x值
        :param df_y: y值，包括y值误差error_y
        :param df_param: 参数信息，name、init、lower、upper
        :param func: 机理函数，入参：x_dict, param_dict, 返回，y值
        :param y_col: df_y中y值对应的列名
        :param error_y_col: df_y中y值误差对应的列名
        :return: 参数拟合结果，列名为name、init、lower、upper、value、error
        """

        # 数据预处理
        self._func = func  # 拟合函数
        x_dict_list: List[Dict[str, float]] = df_x.to_dict("records")  # x列表
        y_list: List[float] = df_y[y_col].astype(float).tolist()  # y真值列表
        error_y_list: List[float] = df_y[error_y_col].astype(float).tolist()  # y误差列表
        param_dict_list: List[Dict[str, str | float]] = df_param.to_dict("records")  # 参数信息列表
        param_dict_dict: Dict[str, Dict[str, str | float]] = {o.get("name"): o for o in param_dict_list}  # 参数信息字典
        param_init_dict: Dict[str, float] = {o.get("name"): o.get("init") for o in param_dict_list}  # 参数初始值字典
        kw_list: List[str] = sorted(list(param_init_dict.keys()))

        # 定义cost函数，最大似然方法中的偏差函数
        def func_cost(*args):
            sum_delta = 0.
            kv_dict = {k: v for k, v in zip(kw_list, args)}
            for x_dict, y, error_y in zip(x_dict_list, y_list, error_y_list):
                y_0 = self._func(x_dict=x_dict, param_dict=kv_dict)
                delta_y = y - y_0
                sum_delta += delta_y * delta_y / (error_y * error_y)
            return sum_delta

        # 初始化最小化器
        m = iminuit.Minuit(func_cost, **param_init_dict, name=kw_list)
        for param_dict in param_dict_list:
            m.limits[param_dict.get("name")] = (param_dict.get("lower"), param_dict.get("upper"))

        # 最小化
        m.migrad()  # find minimum
        m.hesse()  # compute uncertainties

        # 提取参数拟合结果
        param_fitted_list: List[_ParamFitted] = []
        for name, value, err in zip(m.parameters, m.values, m.errors):
            param_dict: Dict[str, float | str] = param_dict_dict.get(name)
            param_fitted_list.append(copy.copy(_ParamFitted(
                name=name,
                lower=param_dict.get("lower"),
                upper=param_dict.get("upper"),
                init=param_dict.get("init"),
                value=value,
                error=err,
            )))
        self._param_fitted_dict = {o.name: o for o in param_fitted_list}
        df_param_fitted = pandas.DataFrame([o.__dict__ for o in param_fitted_list])
        return df_param_fitted

    def _predict_inner(
            self,
            x_dict: Dict[str, float],
            func_error=None,
    ) -> Tuple[float, float | None]:
        """
        推理
        :param x_dict: x值
        :param func_error: 误差计算函数，输入参数 x_dict 为x值的字典、param_dict为参数值字典，输出 float 为y值误差。
        :return: y值和其误差。当误差计算函数为None时，结果的误差值为None。
        """
        y = self._func(x_dict=x_dict, param_dict={k: v.value for k, v in self._param_fitted_dict.items()})

        if func_error is None:  # 当误差函数已知
            return y, None
        else:  # 当误差函数未知
            param_dict_dict = {k: v.__dict__ for k, v in self._param_fitted_dict.items()}
            error_y = func_error(x_dict=x_dict, param_dict=param_dict_dict)
            return y, error_y
